import tensorflow as tf
import util
import matplotlib.pyplot as plt

tf.reset_default_graph()
meta = tf.train.import_meta_graph('./model/model.meta')

sess = tf.Session()
meta.restore(sess, tf.train.latest_checkpoint('./model'))

graph = tf.get_default_graph()
input_x_placeholder = graph.get_tensor_by_name('input_x_placeholder:0')
input_y_placeholder = graph.get_tensor_by_name('input_y_placeholder:0')
predict = graph.get_tensor_by_name('predict:0')
accuracy = graph.get_tensor_by_name('accuracy:0')

test_set_loader = util.DataLoader('test-set')

test_x, test_y = test_set_loader.load_all()
current_accuracy = sess.run(accuracy, feed_dict={
    input_x_placeholder: test_x,
    input_y_placeholder: test_y
})
print('Model prediction accuracy on test set = %f' % current_accuracy)

while True:
    print('> ', end='')
    instruction = input()
    if instruction == 'exit': exit(0)
    try:
        idx = int(instruction)
        if idx >= test_set_loader.total:
            print("Too large an index:", idx)
            continue
    except ValueError as e:
        print("Invalid index format:", instruction)
        continue
    x, y = test_set_loader.load(idx, 1)
    ans = sess.run(predict, feed_dict={
        input_x_placeholder: x,
        input_y_placeholder: y
    })[0]
    print("Prediction result:", ans)
    plt.imshow(x[0][:, :, 0], cmap='gray')
    plt.show()